import argparse
from ultralytics import YOLO
from PIL import Image, ImageDraw
import os


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/train/yolov10_custom/weights/best.pt',
                        help='model.pt path')
    parser.add_argument('--img', type=int, default=640, help='image size')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    return parser.parse_args()


def load_image(source):
    if os.path.isfile(source):
        return [Image.open(source)]
    elif os.path.isdir(source):
        return [Image.open(os.path.join(source, file)) for file in os.listdir(source) if
                file.endswith(('jpg', 'jpeg', 'png'))]
    else:
        raise ValueError("Invalid source path")


def main(opt):
    model = YOLO(opt.weights)
    images = load_image(opt.source)

    for img in images:
        results = model(img, size=opt.img)
        draw = ImageDraw.Draw(img)
        for result in results:
            for box in result.boxes:
                draw.rectangle(box.xyxy, outline="red", width=3)
        img.show()


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
